
#ifndef AFTERMATH_PROBABILITY_DIST_BINOMIAL_HPP_INCLUDED
#define AFTERMATH_PROBABILITY_DIST_BINOMIAL_HPP_INCLUDED

#include "dist_traits.hpp"

#include <cmath>
#include <cstdint>
#include <string>
#include <utility>

namespace aftermath
{
    namespace probability
    {
        // Forward declaration to define traits.
        struct dist_binomial;

        template <>
        struct is_discrete<dist_binomial>
        {
            typedef dist_binomial distribution_type;

            static const bool value = true;
        };

        struct dist_binomial
        {
            typedef dist_binomial type;
            typedef std::size_t result_type;

            static const std::string name;

        private:
            result_type m_number_of_trials;
            double m_probability_of_success;
            double m_cache_probability_of_failure, m_cache_expected_value;

        public:
            dist_binomial()
                : m_number_of_trials(1UL), m_probability_of_success(0.5),
                m_cache_probability_of_failure(0.5), m_cache_expected_value(0.5)
            {
            }

            explicit dist_binomial(result_type n, double probability_of_success)
                : m_number_of_trials(n), m_probability_of_success(probability_of_success),
                m_cache_probability_of_failure(1.0 - probability_of_success), m_cache_expected_value(n * probability_of_success)
            {
            }

            result_type number_of_trials() const { return this->m_number_of_trials; }
            double probability_of_success() const { return this->m_probability_of_success; }
            double probability_of_failure() const { return this->m_cache_probability_of_failure; }

            double expected_value() const { return this->m_cache_expected_value; }
            double variance() const { return this->m_cache_expected_value * this->m_cache_probability_of_failure; }
            double standard_deviation() const { return std::sqrt(this->variance()); }

            double mean() const { return this->expected_value(); }
            double stddev() const { return this->standard_deviation(); }

            double cdf(result_type k) const
            {
                if (k >= this->m_number_of_trials) return 1.0;
                
                double p = 0.0;
                for (result_type j = 0UL; j <= k; j++) p += this->pdf(j);
                return p; 
            }

            double pdf(result_type k) const
            {
                return dist_binomial::n_choose_k(this->m_number_of_trials, k) * std::pow(this->m_probability_of_success, k) * std::pow(this->m_cache_probability_of_failure, this->m_number_of_trials - k);
            }

            static double n_choose_k(result_type n, result_type k)
            {
                if (k > n) return 0.0;
                if (k > n / 2) k = n - k;

                double result = 1.0;
                result_type numerator = n - k;
                for (result_type i = 1UL; i <= k; i++) result *= static_cast<double>(++numerator) / i;
                return result;
            }

            bool operator ==(const type& other) const
            {
                return
                    this->m_number_of_trials == other.m_number_of_trials &&
                    this->m_probability_of_success == other.m_probability_of_success;
            }

            bool operator !=(const type& other) const
            {
                return !this->operator==(other);
            }
        };
        
        const std::string dist_binomial::name = "binomial";
    }
}

namespace std
{
    template <>
    struct hash<aftermath::probability::dist_binomial>
    {
        typedef aftermath::probability::dist_binomial argument_type;
        typedef std::size_t result_type;

        result_type operator()(argument_type const& x) const
        {
            return std::hash<std::size_t>()(x.number_of_trials()) ^ std::hash<double>()(x.probability_of_success());
        }
    };
}

#endif // AFTERMATH_PROBABILITY_DIST_BINOMIAL_HPP_INCLUDED
